import torch
import torch.nn.functional as F
import torch.optim as optim
from torch_geometric.data import Data, Batch
from torch_geometric.nn import GCNConv


def preprocess_data(bom_structures):
    """
    预处理多个BOM结构,构建图数据及标签
    """
    all_components = set()
    edge_lists = []
    label_lists = []

    for bom_structure in bom_structures:
        components = set(part for product, contents in bom_structure.items() for part in contents.keys())
        components.update(
            subpart for product, contents in bom_structure.items() for subparts in contents.values() for subpart in
            subparts)
        all_components.update(components)

        edge_list = []
        label_list = []
        for product, contents in bom_structure.items():
            for part, subparts in contents.items():
                for subpart in subparts:
                    edge_list.append((part, subpart))
                    label_list.append((part, subpart))

        edge_lists.append(edge_list)
        label_lists.append(label_list)

    all_components = sorted(all_components)
    component_to_index = {comp: index for index, comp in enumerate(all_components)}

    graph_datas = []
    labels = []
    for edge_list, label_list in zip(edge_lists, label_lists):
        edge_index = torch.tensor([[component_to_index[edge[0]], component_to_index[edge[1]]] for edge in edge_list],
                                  dtype=torch.long).t().contiguous()
        n_components = len(all_components)
        node_features = torch.eye(n_components)
        graph_data = Data(x=node_features, edge_index=edge_index)
        graph_datas.append(graph_data)

        label = torch.zeros(n_components, n_components, dtype=torch.float)
        for src, dst in label_list:
            label[component_to_index[src], component_to_index[dst]] = 1.0
        labels.append(label)

    return Batch.from_data_list(graph_datas), labels, component_to_index


class GNNModel(torch.nn.Module):
    def __init__(self, num_node_features, hidden_size, out_size):
        super(GNNModel, self).__init__()
        self.conv1 = GCNConv(num_node_features, hidden_size)
        self.conv2 = GCNConv(hidden_size, out_size)

    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x


def train(model, graph_data, labels, optimizer, criterion, epochs=500):
    model.train()
    for epoch in range(epochs):
        optimizer.zero_grad()
        out = model(graph_data.x, graph_data.edge_index)
        loss = 0
        for i, label in enumerate(labels):
            loss += criterion(out[graph_data.batch == i], label)
        loss.backward()
        optimizer.step()
        if epoch % 10 == 0:
            print(f"Epoch {epoch}: Loss {loss.item()}")


def recommend_subcomponents(model, graph_data, parent_component_name, component_to_index, threshold=0.5):
    model.eval()
    with torch.no_grad():
        embeddings = model(graph_data.x, graph_data.edge_index)
        parent_index = component_to_index[parent_component_name]
        parent_embedding = embeddings[parent_index].unsqueeze(0)

        child_indices = graph_data.edge_index[1][graph_data.edge_index[0] == parent_index]
        child_embeddings = embeddings[child_indices]
        similarities = F.cosine_similarity(parent_embedding, child_embeddings)

        sorted_indices = similarities.argsort(descending=True)
        unique_indices = torch.unique(child_indices[sorted_indices])
        components = [list(component_to_index.keys())[i] for i in unique_indices.tolist()]
        scores = [similarities[sorted_indices[i]].item() for i in range(len(unique_indices))]

        recommendations = [(component, score) for component, score in zip(components, scores) if score >= threshold]
        recommendations.sort(key=lambda x: x[1], reverse=True)
        return recommendations


part_properties = {
    "显示屏": {"大小": "10寸", "材质": "玻璃", "重量": "200g"},
    "主控制板": {"型号": "ABC123", "功耗": "50W"},
    "元器件": {"品牌": "XYZ", "尺寸": "5mm"},
    "键盘": {"类型": "机械键盘", "颜色": "黑色", "重量": "300g"},
    "中央处理器": {"型号": "Intel Core i7", "功耗": "65W"}
}
# 示例用法
bom_structure1 = {
    "显示屏": {
        "屏幕": ["玻璃", "元器件"],
        "主控制板": ["元器件", "玻璃"],
        "液晶模组": ["液晶面板", "背光模组"]
    },
    "键盘": {
        "键帽": [],
        "电路板": ["元器件", "金属丝"],
        "主控制板": ["元器件", "玻璃"],
    },
    "中央处理器": {
        "运算器": ["元器件"],
        "控制单元": [{"主控制板": ["元器件", "玻璃"]}, "元器件"],
        "寄存器": ["元器件"]
    }
}
bom_structure2 = {
    "WGPart1": {
        "SBB1": [],
        "SBB2": ["SBB4"],
        "SBB5": ["SBB4", "SBB2"]
    },
    "WGPart2": {
        "SBB3": ["SBB4"],
        "SBB2": ["SBB1"],
        "SBB4": ["Part3"],
    }
}
bom_structure3 = {
    "WGPart1": {
        "SBB1": [],
        "SBB2": ["SBB3", "SBB5"],
        "SBB3": ["SBB6"]
    },
    "WGPart4": {
        "SBB1": [],
        "SBB3": ["SBB5"]
    }
}

graph_data, labels, component_to_index = preprocess_data([bom_structure1, bom_structure2, bom_structure3])

hidden_size = 8
out_size = len(component_to_index)
model = GNNModel(graph_data.num_node_features, hidden_size, out_size)

optimizer = optim.Adam(model.parameters(), lr=0.01)
criterion = torch.nn.BCEWithLogitsLoss()

train(model, graph_data, labels, optimizer, criterion)

parent_component = "SBB3"
recommendations = recommend_subcomponents(model, graph_data, parent_component, component_to_index, threshold=0.6)
print(f"为{parent_component}推荐的子部件及相似度得分: {recommendations}")
